package aiproject3.predictors;

import java.util.TreeSet;

import aiproject3.models.KnowledgeBase;

public class DistributionPredictor<U> implements Predictor<U> {

    private KnowledgeBase<U,Double> knowledgeBase;
    private TreeSet<U> unitsSeen;
    
    public DistributionPredictor(KnowledgeBase<U,Double> kb) {
        knowledgeBase = kb;
        unitsSeen = new TreeSet<U>();
    }
    /**
     * Predicts the next character in the sequence by seeing which Gaussian
     * distribution's pdf yields the most likely unit.
     */
    public U predictNext() {
        
        U best = null;
        double max = -1.0f;
        
        for(U u : unitsSeen) {
            double temp = knowledgeBase.getValueFromModel(u);
            
            //System.out.println(u + "(" + temp + ") > " + best + "(" + max + ")");
            
            if(temp > max) {
                max = temp;
                best = u;
            }
        }
        return best;
    }

    /**
     * Updates the probability distribution of the given unit.
     */
    public void updatePredictor(U unit) {
        
        if(!unitsSeen.contains(unit)) unitsSeen.add(unit);
        
    }

}
